import os
import sys

from vearch.config import Config
from vearch.core.space import Space
from vearch.core.vearch import Vearch
from vearch.filter import Filter, FieldValue, Condition
from vearch.schema.field import Field
from vearch.schema.index import ScalarIndex, HNSWIndex
from vearch.schema.space import SpaceSchema
from vearch.utils import DataType, MetricType, VectorInfo

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

from app.util.embedding import bge_zh_small


# https://github.com/vearch/vearch/blob/master/examples/python/example.py
def create_db(vc: Vearch, db_name: str):
    ret = vc.create_database(db_name)
    print("create db: ", ret.dict_str())
    pass


def create_space(vc: Vearch, db_name: str, space_name: str, dim: int = 512):
    fields = []
    for name in ['source', 'className', 'type']:
        col = Field(name, DataType.STRING, desc=name, index=ScalarIndex(f"{name}_idx"))
        fields.append(col)

    for name in ['safeChat', 'text']:
        col = Field(name, DataType.STRING, desc=name)
        fields.append(col)

    vector = Field(
        "vector",
        DataType.VECTOR,
        HNSWIndex("vector_idx", MetricType.Inner_product, 32, 40),
        dimension=dim,
    )
    fields.append(vector)

    space_schema = SpaceSchema(
        space_name,
        fields=fields,
        description="",
        partition_num=1,
        replica_num=1,
        # 用于生产时设置为3
        # replica_num=3,
    )

    ret = vc.create_space(db_name, space_schema)
    print("create space: ", ret.dict_str())

def drop_space(vc: Vearch, db_name: str, space_name: str):
    ret = vc.drop_space(db_name, space_name)
    print("create space: ", ret.dict_str())


#     {'id': 4, 'name': 'hnsw_bge512_20250826', 'resource_name': 'default', 'version': 2, 'db_id': 3, 'enabled': True, 'partitions': [{'id': 4, 'name': '', 'space_id': 4, 'db_id': 3, 'partition_slot': 0, 'replicas': [1, 2, 3], 'resourceExhausted': False}], 'partition_num': 1, 'replica_num': 3, 'fields': [{'name': 'source', 'type': 'string', 'desc': 'source', 'index': {'name': 'source_idx', 'type': 'SCALAR'}}, {'name': 'className', 'type': 'string', 'desc': 'className', 'index': {'name': 'className_idx', 'type': 'SCALAR'}}, {'name': 'type', 'type': 'string', 'desc': 'type', 'index': {'name': 'type_idx', 'type': 'SCALAR'}}, {'name': 'safeChat', 'type': 'string', 'desc': 'safeChat'}, {'name': 'text', 'type': 'string', 'desc': 'text'}, {'name': 'vector', 'type': 'vector', 'desc': '', 'dimension': 512, 'index': {'name': 'vector_idx', 'type': 'HNSW', 'params': {'metric_type': 'InnerProduct', 'nlinks': 32, 'efConstruction': 40, 'efSearch': 64}}}], 'index': {'name': 'vector_idx', 'type': 'HNSW', 'params': {'metric_type': 'InnerProduct', 'nlinks': 32, 'efConstruction': 40, 'efSearch': 64}}, 'space_properties': {'className': {'field_type': 4, 'type': 'string', 'index': {'name': 'className_idx', 'type': 'SCALAR'}, 'option': 1}, 'safeChat': {'field_type': 4, 'type': 'string'}, 'source': {'field_type': 4, 'type': 'string', 'index': {'name': 'source_idx', 'type': 'SCALAR'}, 'option': 1}, 'text': {'field_type': 4, 'type': 'string'}, 'type': {'field_type': 4, 'type': 'string', 'index': {'name': 'type_idx', 'type': 'SCALAR'}, 'option': 1}, 'vector': {'field_type': 5, 'type': 'vector', 'index': {'name': 'vector_idx', 'type': 'HNSW', 'params': {'metric_type': 'InnerProduct', 'nlinks': 32, 'efConstruction': 40, 'efSearch': 64}}, 'dimension': 512, 'option': 1}}}


def test(vc: Vearch, db_name: str, space_name: str):
    data = []
    book_item = {"_id": "1234",
                 "source": "source",
                 "className": "className",
                 "type": "type",
                 "safeChat": "safeChat",
                 "text": "你好",
                 "vector": bge_zh_small.embed_query('你好')
                 }
    data.append(book_item)
    ret = vc.upsert(db_name, space_name, data)
    print(f'upsert: {ret}')
    document_ids = ret.get_document_ids()
    print(f'upsert: {document_ids}')

    feature = bge_zh_small.embed_query('你好')
    vi = VectorInfo("vector",
                    feature,
                    min_score=-1,
                    max_score=10000
                    )
    conditons_search = [
        Condition(operator="IN", fv=FieldValue("source", ["source"]))
    ]
    search_filters = Filter(operator="AND", conditions=conditons_search)
    ret = vc.search(db_name,
                    space_name,
                    # vector_infos=[vi],
                    vector_infos=[vi],
                    filter= search_filters,
                    vector=True,
                    limit=5,
                    )
    print("search document", ret.documents)
    # search document [[{'_id': '123', '_score': 0.10541725158691406, 'className': 'className', 'safeChat': 'safeChat', 'source': 'source', 'text': 'text', 'type': 'type'}]]

    # ret = vc.delete("database_test", "book_info", ['123'])
    # print(ret.document_ids)
    pass


if __name__ == '__main__':

    master_host = "http://safety-vearch:8817"
    router_host = "http://safety-vearch:9001"
    token = ""

    db_name = "db0"
    collection = "hnsw_bge512_20250826"

    master_config = Config(host=master_host, token=token)
    master_vc = Vearch(master_config)

    create_db(master_vc, db_name=db_name)
    create_space(master_vc, db_name=db_name, space_name=collection)
    # drop_space(vc, db_name=db_name, space_name=space_name)

    router_config = Config(host=router_host, token=token)
    router_vc = Vearch(router_config)
    test(router_vc, db_name, collection)

    pass


# search document [[{'_id': '1234', '_score': 1, 'className': 'className', 'safeChat': 'safeChat', 'source': 'source', 'text': '你好', 'type': 'type', 'vector': [0.050388563, 0.043097533, 0.08578278, 0.006132055, 0.04468228, -0.035420716, 0.025161898, 0.02821077, 0.024695734, 0.027889604, 0.10825354, -0.3072985, -0.044584908, -0.0006573894, -0.022935765, 0.009746429, -0.013807017, 0.008809637, -0.06698811, -0.0309891, 0.009912678, -0.041775867, 0.008844503, -0.0070835575, 0.040529933, 0.027941572, 0.022158071, -0.038005963, -0.05458459, -0.025864545, -0.0031814042, -0.089813136, -0.0055096727, 0.009368922, -0.034590963, -0.019204333, -0.05919954, 0.01399057, -0.01651235, 0.030059563, -0.015221745, 0.007623266, -0.02415061, -0.0061225113, 0.00636114, -0.046605825, -0.06339538, 0.059199873, -0.0330818, 0.041680045, -0.020908471, -0.014462371, 0.04947449, 0.03681892, 0.049069908, -0.05786184, -0.004430908, 0.041538045, -0.045824185, 0.006755462, -0.0068990355, -0.053409938, 0.00915115, 0.0065485276, 0.07340378, -0.0014169967, -0.035384346, -0.013691912, -0.011196887, 0.06581094, 0.020544115, -0.01287655, -0.0130192675, -0.002299029, 0.02346582, -0.060713608, 0.0071689296, 0.005331883, -0.029648917, 0.0059826225, 0.045073137, 0.04094436, 0.009920859, 0.07857029, -0.008965898, -0.0038952935, -0.09045658, -0.012496673, -0.010171273, 0.015415266, -0.11215484, -0.0707497, -0.009713638, 0.04646449, -0.0037354773, 0.06629629, -0.05412792, -0.00082517107, 0.05046202, -0.071611874, -0.01239049, 0.025103664, -0.01625868, 0.010164083, -0.048675038, 0.054094747, -0.07124025, -0.03489638, 0.08365127, -0.0080652535, -0.03773418, -0.054234706, -0.052512996, 0.072169095, -0.032369155, 0.007069063, -0.0017151823, -0.019261155, -0.015089113, 0.09525265, -0.03419814, 0.026638912, 0.051528156, -0.03983782, -0.06695347, -0.040240344, 0.016066115, -0.023309587, -0.074037865, -0.056482513, 0.064410396, 0.053394612, 0.015907241, -0.04337155, -0.015022993, 0.011675462, -0.0011368899, 0.01282373, -0.022957237, 0.00479243, -0.019264383, 0.03235046, -0.09358658, -0.047921248, 0.0085709635, -0.0042479327, 0.009914825, 0.04342511, 0.10224937, 0.044771742, -0.011997223, 0.0100814765, 0.026200125, -0.013175717, -0.07928129, 0.033421066, -0.033058457, 0.06825466, -0.019963417, -0.040432513, 0.07355742, -0.021440916, 0.015285103, 0.09360068, -0.036646664, 0.0004930325, 0.0024474047, 0.06991363, -0.024692945, -0.0072767055, -0.046748534, -0.01882289, -0.10496534, 0.007935826, -0.0034714614, -0.044441853, 0.004171532, -0.08252438, 0.03513616, -0.005135873, -0.057433397, -0.021391736, -0.020810314, -0.031918276, -0.021958774, 0.008851164, -0.010881295, -0.027042026, -0.02270528, -0.012919095, 0.0143121835, -0.016587522, 0.017444769, -0.036757786, -0.018504579, -0.027139252, -0.010013442, 0.038369525, -0.017097905, -0.01840351, -0.016755357, -0.010035525, 0.018314635, 0.016545631, 0.020981636, -0.030555567, -0.02163946, -0.014386795, -0.034738854, -0.012086864, -0.018019458, -0.02261126, -0.030231485, -0.0022701875, 0.037699264, 0.03688395, 0.043760445, -0.0025728869, 0.010884951, -0.05229569, 0.017747557, 0.017748207, 0.08319627, -0.081579715, -0.028087296, 0.044624537, -0.015378538, 0.033493347, -0.023586087, -0.038164314, 0.032516472, 0.06068928, -0.019521369, -0.050404448, 0.07497857, -0.012499739, -0.016916936, 0.027757913, 0.018317565, 0.010857676, -0.0667614, 0.13955703, -0.027709994, 0.048831046, -0.04116368, -0.021014124, -0.033128493, -0.014005943, 0.0007959374, -0.06729309, -0.023342429, -0.005743864, -0.045640323, -0.014219949, 0.037608396, -0.03984148, 0.027665157, 0.0186069, 0.013397582, -0.008726777, -0.032945614, 0.040764485, -0.066969596, 0.040239926, 0.027825488, -0.006117559, 0.013653183, 0.054689012, -0.0801198, -0.050540116, -0.0069896686, -0.031828817, -0.010739878, -0.018996887, -0.040264953, -0.059153177, -0.024578491, 0.037475906, 0.032925602, 0.010166252, 0.034166217, 0.0035694288, 0.020238027, 0.05664536, -0.081259795, 0.020882994, -0.066092506, 0.02216489, 0.0024390607, -0.014864271, -0.0055030193, 0.0063351765, 0.05967161, 0.039650407, -0.007293486, -0.03606442, -0.03440762, -0.0152740665, 0.017371237, -0.010807436, 0.010229482, 0.016729552, -0.022194058, -0.023602998, -0.029992726, -0.011967815, -0.023971712, -0.016660463, 0.027467726, -0.0011834, -0.04749478, -0.025937537, 0.011776862, -0.045235824, -0.018534832, 0.0040826364, 0.07831619, -0.05383237, -0.02273101, -0.04176761, -0.019811496, 0.02536425, 0.00058855815, -0.0015746539, 0.05283377, 0.028013319, 0.062406503, 0.02402538, -0.051996846, 0.01752146, 0.024328327, -0.003774579, 0.018935688, -0.006874256, -0.06380716, -0.022424772, -0.012755234, -0.008277752, -0.0006159908, 0.019278128, 0.017668, 0.04252601, -0.006710298, 0.045697305, 0.004900863, 0.026732812, -0.0022969719, -0.055616587, 0.029260645, 0.010884033, 0.0152434325, -0.0070363316, 0.032424644, -0.06328252, -0.003102844, -0.043394666, 0.07474792, -0.0074603376, 0.010228133, 0.054784354, -0.060690504, -0.035236023, -0.023676729, -0.020124722, 0.06556535, -0.012475404, -0.012718512, -0.013404157, -0.015781367, 0.008244492, -0.014761105, -0.0027441755, 0.06429653, -0.075261235, 0.059062414, -0.071049534, 0.012038969, -0.03436796, -0.044394344, -0.053434554, 0.030759793, 0.056089435, -0.022733903, -0.039828412, 0.03725507, -0.08115502, -0.085108094, 0.012936914, 0.05287962, 0.008458963, 0.03928006, -0.022429354, 0.06666703, 0.022971027, 0.029526401, -0.096903235, -2.1202888e-05, -0.009081124, 0.052801143, -0.04230996, 0.055446688, -0.016868811, -0.018356344, 0.0022048217, -0.0015982861, 0.018374566, 0.020803086, 0.04138737, 0.027354363, -0.06662121, 0.25313216, -0.05880687, -0.012670449, 0.041599255, 0.029473498, 0.0013140304, -0.07115439, 0.03707719, 0.020557968, -0.02225671, -0.06963824, -0.019772947, 0.011533118, -0.051747847, 0.080427654, -0.021497514, -0.04116022, -0.059837278, -0.057160567, -0.03188131, 0.05515405, 0.025047602, 0.010125998, -0.09685871, -0.0067910054, 0.0018720608, -0.044154737, 0.056610737, -0.0006150504, -0.034163818, -0.031356502, 0.026966348, 0.012882855, -0.0011213084, 0.04067553, 0.070761286, 0.057963137, 0.0510743, 0.027808579, -0.026982488, 0.029014062, 0.039454564, 0.0063996273, 0.04357906, 0.035153948, 0.07629337, -0.06337824, 0.025268609, 0.04058105, -0.08769158, -0.017534992, -0.05743738, -0.014070471, 0.03322317, -0.063494906, -0.005673951, 0.0013136141, -0.0032616546, -0.0129691865, 0.02484209, -0.05617286, -0.024516845, -0.04958403, -0.003426862, 0.026348468, 0.076764256, -0.026788324, 0.086502135, 0.062554605, 0.02897422, -0.021507764, -0.07418549, 0.015656127, 0.0694363, 0.0003475432, 0.096514955, 0.043549847, -0.036487512, 0.053491555, -0.016332172, -0.058854133, 0.021626994, -0.0018287532, -0.01483157, -0.048685715, -0.019347224, -0.014440354, 0.050009463, -0.009526756, -0.028338324, -0.02857364, 0.034640186, -0.021733567, -0.02162404, -0.071297705, -0.019214239, 0.058321506, 0.026304524, -0.05874354, 0.030416062, 0.08117939, 0.0022759826]}, {'_id': '123', '_score': 0.10541725158691406, 'className': 'className', 'safeChat': 'safeChat', 'source': 'source', 'text': 'text', 'type': 'type', 'vector': [0.5399088, 0.635592, 0.46142635, 0.02551171, 0.78367376, 0.63332814, 0.45851886, 0.68838304, 0.34537226, 0.63662225, 0.9312115, 0.10743668, 0.26675647, 0.59058565, 0.4784452, 0.17386582, 0.60930645, 0.012249579, 0.56284267, 0.3883514, 0.8097924, 0.83440405, 0.71553236, 0.9618067, 0.99791616, 0.5526203, 0.48210332, 0.74447006, 0.12150801, 0.098751575, 0.9666872, 0.35806012, 0.40511116, 0.45195717, 0.12506305, 0.10021336, 0.2942703, 0.89217275, 0.97666615, 0.20556803, 0.32166305, 0.9226453, 0.6343405, 0.20529997, 0.54327065, 0.41879454, 0.7991946, 0.21029352, 0.17797469, 0.057540976, 0.042370815, 0.9836173, 0.16690469, 0.25837222, 0.52002287, 0.15104964, 0.70433915, 0.34456155, 0.7229796, 0.3110323, 0.030180598, 0.8072619, 0.8206055, 0.3125882, 0.1259588, 0.5841295, 0.36238557, 0.17192581, 0.6734517, 0.59477806, 0.0036408245, 0.86931646, 0.25602028, 0.6764919, 0.7915204, 0.29203746, 0.46889925, 0.69752187, 0.2872905, 0.18098165, 0.863004, 0.11589905, 0.837338, 0.11882273, 0.4493001, 0.38792452, 0.5448492, 0.2644258, 0.06644779, 0.78761345, 0.53516954, 0.60774446, 0.5357729, 0.55303437, 0.39469665, 0.42268708, 0.18806256, 0.38599023, 0.571431, 0.0064454987, 0.35501277, 0.64544123, 0.15156855, 0.20755273, 0.54621017, 0.99643755, 0.71260214, 0.8365983, 0.41095182, 0.929023, 0.41829643, 0.32519218, 0.5524959, 0.27208757, 0.70651174, 0.3645437, 0.35600722, 0.5702221, 0.9658382, 0.2633607, 0.20197664, 0.21225692, 0.4694944, 0.015954928, 0.71234286, 0.038488913, 0.5830201, 0.034958724, 0.7185988, 0.6720302, 0.36136267, 0.22294821, 0.7752616, 0.7099867, 0.17751768, 0.03775598, 0.4518907, 0.14595367, 0.5513031, 0.12725656, 0.619694, 0.40585178, 0.2466114, 0.4885457, 0.907543, 0.16954428, 0.84472597, 0.5061228, 0.35106137, 0.10970158, 0.34671146, 0.0527479, 0.99320334, 0.9542778, 0.11034771, 0.30485326, 0.47072235, 0.6429008, 0.38494778, 0.6980975, 0.44052795, 0.3975721, 0.26271662, 0.6322282, 0.039495718, 0.61372864, 0.42378256, 0.76286846, 0.15809679, 0.88077396, 0.07725931, 0.5067348, 0.04330204, 0.28240815, 0.7410381, 0.08009217, 0.3686371, 0.48335472, 0.4196303, 0.95911354, 0.45093426, 0.73557574, 0.04220591, 0.74757206, 0.16247857, 0.80659807, 0.548928, 0.87352175, 0.09059632, 0.7510746, 0.8659352, 0.4279928, 0.28299874, 0.42503795, 0.95390195, 0.6170972, 0.6941124, 0.28904462, 0.095644034, 0.41767833, 0.13530073, 0.31223056, 0.20851272, 0.15898784, 0.38904884, 0.6111747, 0.29266086, 0.864577, 0.9274977, 0.83116376, 0.53461456, 0.7072055, 0.14096634, 0.749673, 0.8391137, 0.8440276, 0.9404132, 0.016719494, 0.2970812, 0.8272623, 0.83242846, 0.4943797, 0.21809614, 0.47760555, 0.8773055, 0.6906562, 0.63892406, 0.025228428, 0.8973531, 0.01538639, 0.5986778, 0.62377465, 0.87300265, 0.48811892, 0.27723837, 0.30223027, 0.2549698, 0.1351825, 0.03883938, 0.27159116, 0.050214857, 0.66137695, 0.48666036, 0.54612184, 0.33648333, 0.80574495, 0.4954525, 0.6450667, 0.8332615, 0.5203608, 0.45499977, 0.7005954, 0.49829155, 0.42255345, 0.8908122, 0.043950416, 0.25980034, 0.97882205, 0.18644166, 0.8735943, 0.35724857, 0.19536828, 0.19266936, 0.64740026, 0.41232723, 0.9300212, 0.26773083, 0.8357974, 0.28891093, 0.48880473, 0.18812893, 0.9322075, 0.64646775, 0.40874952, 0.97665656, 0.6873658, 0.8261439, 0.19402586, 0.9103462, 0.85205597, 0.5296492, 0.5459738, 0.5589138, 0.60978925, 0.1744622, 0.37673345, 0.54879874, 0.75686026, 0.26452368, 0.98422086, 0.36266136, 0.4462028, 0.35821524, 0.09113442, 0.32529646, 0.23043604, 0.8250085, 0.17514375, 0.3490241, 0.86065716, 0.5072521, 0.20322286, 0.96603584, 0.4638194, 0.65079874, 0.44081107, 0.8148, 0.13474514, 0.6720728, 0.574922, 0.39800456, 0.79235977, 0.9275379, 0.24665186, 0.10970034, 0.46843985, 0.5019835, 0.30599344, 0.058271933, 0.23364331, 0.94840884, 0.36327136, 0.93621045, 0.09013136, 0.53243196, 0.14248623, 0.5368431, 0.7426848, 0.3300642, 0.4833603, 0.36910155, 0.7270526, 0.09188352, 0.6461948, 0.9672123, 0.0300598, 0.31859002, 0.7821508, 0.6416717, 0.016535742, 0.31331795, 0.0076415455, 0.009925281, 0.93756014, 0.7582613, 0.65473354, 0.4726909, 0.22757232, 0.8903621, 0.69391245, 0.9263544, 0.3871695, 0.5505129, 0.37939495, 0.5173802, 0.46147892, 0.6843344, 0.12533347, 0.9066603, 0.43420112, 0.0658545, 0.55749255, 0.8710934, 0.7072118, 0.64828926, 0.4313525, 0.08356771, 0.22777444, 0.19018753, 0.9244094, 0.07531478, 0.42696473, 0.25438657, 0.038775776, 0.41450527, 0.45635644, 0.22905463, 0.5682802, 0.13212763, 0.9064798, 0.18124694, 0.6582996, 0.65151674, 0.52902305, 0.6682778, 0.21789606, 0.27809182, 0.7881893, 0.3744722, 0.69434214, 0.3969867, 0.5394561, 0.8628981, 0.11963688, 0.2283156, 0.6983952, 0.61412054, 0.62296516, 0.38505918, 0.035098698, 0.2714104, 0.839457, 0.64094794, 0.5857927, 0.5492963, 0.6117619, 0.7943526, 0.039733704, 0.4939915, 0.036092445, 0.9938507, 0.17075035, 0.45882633, 0.15846935, 0.80306727, 0.5604522, 0.8670019, 0.6070375, 0.5756687, 0.33984387, 0.5650219, 0.4080838, 0.6029838, 0.16531512, 0.38447225, 0.6561874, 0.8201837, 0.08174848, 0.43402395, 0.44480026, 0.3069339, 0.3639857, 0.073540404, 0.27491686, 0.36169407, 0.19937775, 0.74315935, 0.82414126, 0.8667751, 0.50417507, 0.085448466, 0.03176122, 0.5592265, 0.7896004, 0.23128591, 0.56549877, 0.78334975, 0.17926732, 0.52759916, 0.0506943, 0.3443799, 0.2604336, 0.3340358, 0.14796102, 0.53306943, 0.93256944, 0.8318016, 0.8250402, 0.0074624065, 0.5589844, 0.91235644, 0.45390233, 0.5710046, 0.49613786, 0.3914853, 0.010897279, 0.79620916, 0.9308112, 0.09492316, 0.20713909, 0.77263683, 0.361299, 0.21679671, 0.5505393, 0.11829267, 0.6327681, 0.24689054, 0.79829633, 3.9158054e-05, 0.25148022, 0.86323565, 0.18512765, 0.3058041, 0.6536517, 0.91483366, 0.35152668, 0.75000745, 0.42214763, 0.7326616, 0.8245618, 0.4344011, 0.8905335, 0.90464073, 0.30829597, 0.15381908, 0.7087119, 0.9803056, 0.5721229, 0.03021249, 0.16119654, 0.81459093, 0.7122838, 0.5036157, 0.22734527, 0.32001582, 0.16047099, 0.8763619, 0.3167067, 0.47640383, 0.5171337, 0.15843585, 0.7018068]}]]
